package com.leiyk.rateLimiter.aop;

import com.leiyk.rateLimiter.annotation.SimpleRateLimiter;
import com.leiyk.rateLimiter.annotation.SlideRateLimiter;
import com.leiyk.rateLimiter.exception.RateLimiterException;
import jakarta.annotation.Resource;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.time.LocalDate;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * @Description: 滑动窗口限流注解aop实现
 * @Author: leiyk6
 * @Date: 2024/7/5 9:54
 */
@Aspect
@Component
public class SlideRateLimiterAop {
    @Pointcut("@annotation(com.leiyk.rateLimiter.annotation.SlideRateLimiter)")
    void pc(){}

    @Resource
    StringRedisTemplate stringRedisTemplate;


    String scr = "local key = KEYS[1]\n" +
            "local now = tonumber(ARGV[1])\n" +
            "local window_size = tonumber(ARGV[2])\n" +
            "local limit = tonumber(ARGV[3])\n" +
            "redis.call('ZREMRANGEBYSCORE', key, 0, now - window_size)\n" +
            "local count = redis.call('ZCARD', key)\n" +
            "if count < limit then\n" +
            "    redis.call('ZADD', key, now, now)\n" +
            "    return 1\n" +
            "else\n" +
            "    return 0\n" +
            "end";


    @Around(value = "pc()")
    Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        SlideRateLimiter annotation = method.getAnnotation(SlideRateLimiter.class);
        int count = annotation.count();
        TimeUnit timeUnit = annotation.timeUnit();
        long windowSize = annotation.windowSize();
        String key = signature.getMethod().toString();
        long realWindowSize = timeUnit.toMillis(windowSize);

        Long res = stringRedisTemplate.execute(new DefaultRedisScript<Long>(scr), List.of(key), String.valueOf(System.currentTimeMillis()), String.valueOf(realWindowSize));
        if(res.equals(0L)){
            throw RateLimiterException.rejectRequest();
        }
        return joinPoint.proceed();
    }

}
